######################################################################
#
# map_lnl_landscape.py
#
# Victor Hanson-Smith
# victorhs@cs.uoregon.edu
#
# OVERVIEW: Creates a 2-D PDF plot of the phylogenetic likelihood
# landscape for two parameters.
#
# This script was created in July 2010 for my research on the separability (or rather,
# the non-separability) of phylogenetic model parameters.
#
######################################################################
import os
import re
import sys

from argParser import ArgParser
argp = ArgParser(sys.argv)

sys.path.append("./")
from map_lnl_setup import *

def invoke_phyml(fcommands, command):
    # invoke PHYML
    print command
    fcommands.write(command + "\n")
    fcommands.flush()
    os.system(command)    

# returns lnl value, parsed from stats file
def parse_phyml_stats(stats_path):
    # parse the stats file from PHYML
    fin = open(stats_path, "r")
    lines = fin.readlines()
    fin.close()
    for line in lines:
        if line.startswith(". Log-likelihood:"):
            tokens = line.strip().split()
            lnl = tokens[2]
            lnl = re.sub("\.$", "", lnl)
    return lnl

def calculate_map():    
    fout = open("./" + POINTS_PATH, "w")
    fcommands = open("command_log.txt", "w")
    
    i = xmin
    while i < xmax:
        
        j = ymin
        if dimensions == 2: # are we exploring two dimensions?
            while j < ymax:
                print i, j
                command = write_phyml_command(i, j) 
                # write a fresh Newick tree (potentially containing incremented BLs)
                treepath = write_tree(i, j, i.__str__() + "." + j.__str__())
                command += " -u " + treepath
        
                invoke_phyml(fcommands, command)
                lnl = parse_phyml_stats(alignment_path + "_phyml_stats.txt")
    
                # write the new i,j,lnl value to our POINTS_PATH
                out_line = i.__str__() + "\t" + j.__str__() + "\t" + lnl.__str__()
                fout.write(out_line + "\n")
                fout.flush()
                
                # increment and loop
                j += step
        elif dimensions == 1:
            print i
            command = write_phyml_command(i, j) 
            # write a fresh Newick tree (potentially containing incremented BLs)
            treepath = write_tree(i, j, i.__str__())
            command += " -u " + treepath
    
            invoke_phyml(fcommands, command)
            lnl = parse_phyml_stats(alignment_path + "_phyml_stats.txt")

            # write the new i,j,lnl value to our POINTS_PATH
            out_line = i.__str__() + "\t" + lnl.__str__()
            fout.write(out_line + "\n")
            fout.flush()
                      
        i += step

    fout.close()
    fcommands.close()

    

def draw_heatmap():    
    # parse the points file:
    fin = open("./" + POINTS_PATH, "r")
    lines = fin.readlines()
    fin.close()
    points = {} # an array of arrays: [x, y, lnl]
    for line in lines:
        line = line.strip()
        tokens = line.split()
        x = tokens[0]
        y = tokens[1]
        lnl = tokens[2]
        if lnl == "-inf":
            lnl = -2**32 # the maximally smallest 32-bit number
        if points.keys().__contains__(x):
            points[x][y] = lnl
        else:
            points[x] = {}
            points[x][y] = lnl
        
    # a CSV file
    fout = open("./" + POINTS_PATH + ".table", "w")
    # write the header line
    ypoints = points[ points.keys()[0] ].keys()
    ypoints.sort()
    line = ','.join( ypoints )
    fout.write("X," + line + "\n")
    xpoints = points.keys()
    xpoints.sort()
    for x in xpoints:
        print x
        line_points = []
        for y in ypoints:
            line_points.append(points[x][y].__str__() )
        line = ','.join( line_points )
        fout.write(x.__str__() + "," + line + "\n")
    fout.close()
    
    # build a CRAN script
    fout = open("./" + CRAN_PATH, "w")
    fout.write("t <- read.csv(\"points.txt.table\", sep=\",\")\n")
    fout.write("row.names(t) <- t$X\n")
    fout.write("t <- t[,2:" + points[ points.keys()[0] ].keys().__len__().__str__() + "]\n")
    fout.write("tm <- data.matrix(t)\n")
    fout.write("pdf(\"heatmap.pdf\")\n")
    # note: xlab and ylab are switched in the following statement because R likes to rotate plots in a confusing way.
    fout.write("h <- heatmap(tm, Rowv=NA, Colv=NA, col=heat.colors(2056), scale=\"column\", margins=c(5,10), ylab=\"" + xlab + "\", xlab=\"" + ylab + "\" )\n")
    fout.write("dev.off()\n")
    fout.close()
    os.system(CRAN + " --save < " + "./" + CRAN_PATH)


def draw_plot():    
    # parse the points file:
    fin = open("./" + POINTS_PATH, "r")
    lines = fin.readlines()
    fin.close()
    points = {} # an array of doubles: [x, lnl]
    for line in lines:
        line = line.strip()
        tokens = line.split()
        x = float(tokens[0])
        lnl = tokens[1]
        if lnl == "-inf":
            lnl = -2**32 # the maximally smallest 32-bit number
        points[x] = float(lnl)
        
    fout = open("./" + CRAN_PATH, "w")
    lstr = "x <- c("
    keys = points.keys()
    keys.sort()
    for x in keys:
        this_x = "%.2f"%x
        lstr += this_x + ","
    lstr = re.sub(",$", "", lstr)
    lstr += ")"
    fout.write(lstr + "\n")
    
    lstr = "lnl <- c("
    for x in keys:
        this_lnl = "%.2f"%points[x]
        lstr += this_lnl + ","
    lstr = re.sub(",$", "", lstr)
    lstr += ")"
    fout.write(lstr + "\n")
    
    fout.write("pdf(\"plot.pdf\")\n")
    fout.write("plot(x, lnl, type=\"b\", pch=20, xlab=\"X\", ylab=\"lnL\")\n")
    fout.write("dev.off()\n")
    fout.close()
    os.system(CRAN + " --save < " + "./" + CRAN_PATH)

########################
#
# main:
#
action = argp.getArg("--action")
if action == "calculate":
    calculate_map()
elif (action == "draw" and dimensions == 2):
    draw_heatmap()
elif (action == "draw" and dimensions == 1):
    print "hello!"
    draw_plot()